import torch
from PIL import Image
from torch import Tensor
from torch.utils.data import Dataset
from torchvision.transforms import Compose
import torch
import matplotlib.pyplot as plt
def convert_to_one_hot(mask):
    # Assuming mask is already an integer tensor with values {0, 1, 2}
    one_hot_mask = torch.nn.functional.one_hot(mask, num_classes=3).permute(2, 0, 1).float()
    return one_hot_mask


def plot_one_hot_channels(one_hot_mask):
    """
    Plots each channel of a one-hot encoded mask as a separate grayscale image.

    :param one_hot_mask: torch.Tensor, shape [C, H, W], one-hot encoded mask
    """
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))  # 1 row, 3 columns of subplots

    # Assuming one_hot_mask has shape [3, H, W] and 3 channels
    for i in range(one_hot_mask.shape[0]):
        ax = axes[i]
        ax.imshow(one_hot_mask[i].numpy(), cmap='gray')
        ax.set_title(f'Channel {i}')
        ax.axis('off')  # Turn off axis numbers and ticks

    plt.show()

class SegmentationDataset(Dataset):
    """A custom dataset class"""

    def __init__(self, images: list, masks: list, transforms: Compose) -> None:
        # Get paths to images and masks
        self.images = images
        self.masks = masks
        self.transforms = transforms

    def __len__(self) -> int:
        # Return the number of images
        return len(self.images)

    def __getitem__(self, index: int) -> (Tensor, Tensor):
        # Get image and mask paths
        image_path = self.images[index]
        mask_path = self.masks[index]

        # Open image and convert to RGB
        image = Image.open(image_path).convert('RGB')
        # Open mask and convert to grayscale
        mask = Image.open(mask_path).convert('L')

        # Apply transforms
        if self.transforms is not None:
            image = self.transforms(image)
            mask = self.transforms(mask)
            # Converting float values to integers
            mask = mask * 255
            mask = mask.squeeze().to(torch.int64)
            # Ground truth labels are 1, 2, 3. therefore subtract one to achieve 0, 1, 2:
            mask -= 1
        mask_new = convert_to_one_hot(mask)
        # plot_one_hot_channels(one_hot_mask)
        # mask_new = mask.unsqueeze(0)
        # Return the transformed image and one-hot encoded mask
        return image, mask_new
        # Return the image and corresponding mask
